package fr.cephb.dbsnptools;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;
import java.io.RandomAccessFile;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;

import com.sleepycat.bind.tuple.BooleanBinding;
import com.sleepycat.bind.tuple.IntegerBinding;
import com.sleepycat.je.DatabaseConfig;
import com.sleepycat.je.Environment;
import com.sleepycat.je.EnvironmentConfig;



import fr.cephb.berkeleydb.db.PrimaryDB;
import fr.cephb.berkeleydb.db.Walker;
import fr.cephb.io.IOUtils;
import fr.cephb.util.AbstractCmdLine;
import fr.cephb.util.FQuickSort;


public class DigestBowtie2
	extends AbstractCmdLine
	{
	/** pattern for a simple mutation */
	private final static  Pattern TAB=Pattern.compile("[\t]");
	/** max size of a Hit as binary */
	private static final int HIT_SIZEOF=38;
	/** suffix for database */
	private static final String SUFFIX=".bindb";
	private static final String SUFFIX_L=".left"+SUFFIX;
	private static final String SUFFIX_R=".right"+SUFFIX;
	private Environment environment;
	private PrimaryDB<Integer, Boolean> rsidSetDb;
	private boolean strict_use_matches=true;
	//private int min_rs_id=Integer.MAX_VALUE;
	//private int max_rs_id=-1;
	private File base=null;
	private String program=null;
	private String ids_file=null;
	
	private static class Hit
		{
		int rs_id;
		char strand;
		int chrom_id;
		int position;
		int length;
		String  mismatch="";
		
		int countMismatches()
			{
			int c=0;
			for(int i=0;i< this.mismatch.length();++i)
				{
				if(this.mismatch.charAt(i)==',') c++;
				}
			return c;
			}

		@Override
		public int hashCode()
			{
			final int prime = 31;
			int result = 1;
			result = prime * result + chrom_id;
			result = prime * result + length;
			result = prime * result + position;
			result = prime * result + rs_id;
			result = prime * result + strand;
			return result;
			}

		@Override
		public boolean equals(Object obj)
			{
			if (this == obj) return true;
			if (obj == null) return false;
			if (getClass() != obj.getClass()) return false;
			Hit other = (Hit) obj;
			if (chrom_id != other.chrom_id) return false;
			if (length != other.length) return false;
			if (position != other.position) return false;
			if (rs_id != other.rs_id) return false;
			if (strand != other.strand) return false;
			return true;
			}
		

		}
	
	private static class CandidatePair
	implements Comparable<CandidatePair>
		{
		Hit L;
		Hit R;
		
		CandidatePair(Hit L,Hit R)
			{
			this.R=R;
			this.L=L;
			}
		
		public int countMismatches()
			{
			return L.countMismatches()+R.countMismatches();
			}
		
		@Override
		public boolean equals(Object o)
			{
			if(o==this) return true;
			if(o==null || !(o instanceof CandidatePair)) return false;
			CandidatePair cp=CandidatePair.class.cast(o);
			return R.equals(cp.R) && L.equals(cp.L);
			}
		
		@Override
		public int hashCode()
			{
			return R.hashCode()+31*L.hashCode();
			}
		
		@Override
		public int compareTo(CandidatePair o)
			{
			return countMismatches()-o.countMismatches();
			}
		@Override
		public String toString()
			{
			return " L :"+L+"\n R :"+R;
			}
		public int getRsId()
			{
			if(R.rs_id!=L.rs_id) throw new RuntimeException("getRsId:"+R.rs_id+"/"+L.rs_id);
			return this.R.rs_id;
			}
		
		public int getChromId()
			{
			if(R.chrom_id!=L.chrom_id) throw new RuntimeException("getChromId:"+toString());
			return this.R.chrom_id;
			}
		
		public char getStrand()
			{
			assert(R.strand==L.strand);
			return this.R.strand;
			}
		
		public int getSNPStart()
			{
			return getStrand()=='+'
				?L.position+ L.length
				:R.position+ R.length
				;
				
			}
		
		public int getSNPEnd()
			{
			return getStrand()=='+'
				?R.position
				:L.position
				;
				
			}
		void print(PrintStream out)
			{
			out.print("rs"+getRsId()+"\t"+
					int2chrom(getChromId())+"\t"+
					getStrand()+"\t"+
					getSNPStart()+"\t"+
					getSNPEnd()
					);
			}
		}
	
	private static class HitComparator
		implements Comparator<byte[]>
		{
		@Override
		public int compare(byte[] o1, byte[] o2)
			{
			try
				{
				ByteArrayInputStream bais1=new ByteArrayInputStream(o1);
				DataInputStream in=new DataInputStream(bais1);
				int rs_id1=in.readInt();
				in.close();
				
				bais1=new ByteArrayInputStream(o2);
				in=new DataInputStream(bais1);
				int rs_id2=in.readInt();
				in.close();
				return rs_id1-rs_id2;
				}
			catch (IOException e)
				{
				throw new RuntimeException(e);
				}
			}
		}
	
	private static class HitBinding
		{
		public Hit read(byte array[]) throws IOException
			{
			ByteArrayInputStream bais=new ByteArrayInputStream(array);
			DataInputStream in=new DataInputStream(bais);
			Hit h= new Hit();
			h.rs_id=in.readInt();
			h.strand=(char)in.readByte();
			h.chrom_id=in.readShort();
			h.position=in.readInt();
			h.length=in.readShort();
			int n= in.readShort();
			byte ss[]=new byte[n];
			in.readFully(ss);
			h.mismatch=new String(ss);
			in.close();
			return h;
			}
		
		public void write(Hit h,byte array[]) throws IOException
			{
			ByteArrayOutputStream baos=new ByteArrayOutputStream(HIT_SIZEOF);
			DataOutputStream out= new DataOutputStream(baos);
			out.writeInt(h.rs_id);
			out.writeByte(h.strand);
			out.writeShort(h.chrom_id);
			out.writeInt(h.position);
			out.writeShort(h.length);
			byte ss[]=h.mismatch.getBytes();
			out.writeShort(ss.length);
			out.write(ss);
			out.flush();
			byte buff[]=baos.toByteArray();
			if(buff.length> array.length) throw new IOException("sizeof out of range : data requires "+buff.length);
			System.arraycopy(buff, 0, array, 0, buff.length);
			}
		}
	
	
	
	private static class HitDatabase
		{
		private File file;
		private FileOutputStream out=null;
		private RandomAccessFile in=null;
		
		public HitDatabase() throws IOException
			{
			
			}
		
		public HitDatabase(File dir,String suffix) throws IOException
			{
			this.file=File.createTempFile("_sort",suffix,dir);
			}
		public void sort() throws IOException
			{
			//MergeSort sort=new MergeSort(HIT_SIZEOF,new HitComparator());
			FQuickSort sort=new FQuickSort(HIT_SIZEOF, new HitComparator());
			//sort.setBufferCapacity(20*1000*1000);//20Mo
			sort.sort(this.file);
			}
		
		public void close()
			{
			
			}
		}
	

	private Map<Integer,HitDatabase> chrom2left=new HashMap<Integer, HitDatabase>();
	private Map<Integer,HitDatabase> chrom2right=new HashMap<Integer, HitDatabase>();
	private int max_len_deletion=100;
	private int expect_flanking_length=120;
	
	private boolean readFully(RandomAccessFile in,byte array[])
	throws IOException
		{
		int nRead=0;
		int n;
		while((n=in.read(array,nRead,array.length-nRead))!=-1)
			{
			nRead+=n;
			if(nRead==array.length) break;
			}
		if(nRead==0) return false;
		if(nRead!=array.length) throw new IOException("Failed to read "+array.length);
		return true;
		}

	
	/** C+ lower_bound */
	private  long lower_bound(
				RandomAccessFile file,
				long first, long last,
	            int rs_id
	            ) throws IOException
	    {
	    byte x[]=new byte[HIT_SIZEOF];
	    HitBinding binding=new HitBinding();
	    long len = last - first;
	    while (len > 0)
	            {
	            long half = len / 2;
	            long middle = first + half;
	            file.seek(middle*HIT_SIZEOF);
	            if(!readFully(file, x)) throw new IOException("cannot read!");
	            Hit hit=binding.read(x);
	            if (hit.rs_id- rs_id<0)
	                    {
	                    first = middle + 1;
	                    len -= half + 1;
	                    }
	            else
	                    {
	                    len = half;
	                    }
	            }
	    return first;
	    }
	
	/** C+ lower_bound */
	private  List<Hit> equal_range(
				RandomAccessFile file,
	            int rs_id
	            ) throws IOException
		{
		HitBinding binding=new HitBinding();
		List<Hit> array=new ArrayList<Hit>();
		long offset= lower_bound(file,0L,file.length()/HIT_SIZEOF,rs_id);
		//System.err.println("Seek "+offset);
		file.seek(offset*HIT_SIZEOF);
		byte x[]=new byte[HIT_SIZEOF];
		while(readFully(file, x))
			{
			Hit h=binding.read(x);
			if(h.rs_id!=rs_id) break;
			//System.err.println(""+rs_id+" "+h.toString()+" "+array.size());
			array.add(h);
			}
		return array;
		}
	
	
	private void open() throws IOException
		{
		EnvironmentConfig envConfig= new EnvironmentConfig();
		DatabaseConfig dbcfg=new DatabaseConfig();
		envConfig.setConfigParam(EnvironmentConfig.LOG_FILE_MAX,"500000000"/* 500 Mo */);
		if(this.program.equals("sort"))
			{
			envConfig.setAllowCreate(true);
			envConfig.setReadOnly(false);
			dbcfg.setAllowCreate(true);
			dbcfg.setReadOnly(false);
			for(File f: this.base.listFiles())
				{
				if(f.getName().endsWith(SUFFIX))
					{
					throw new IOException("Error "+this.base+" should be empty but contains already "+f);
					}
				}
			
			}
		else if(this.program.equals("dump"))
			{
			dbcfg.setAllowCreate(false);
			dbcfg.setReadOnly(true);
			envConfig.setAllowCreate(false);
			envConfig.setReadOnly(true);
			int index=0;
			for(File f: this.base.listFiles())
				{
				if(f.getName().endsWith(SUFFIX_L))
					{
					HitDatabase db=new HitDatabase();
					db.file=f;
					this.chrom2left.put(index++, db);
					}
				else if(f.getName().endsWith(SUFFIX_R))
					{
					HitDatabase db=new HitDatabase();
					db.file=f;
					this.chrom2right.put(index++, db);
					}
				}
			}
		this.environment= new Environment(this.base, envConfig);
		this.rsidSetDb=new PrimaryDB<Integer, Boolean>(this.environment, null,"snpids",dbcfg, new IntegerBinding(), new BooleanBinding());
		}
	
	private void close()
		{
		for(HitDatabase db:this.chrom2left.values())
			{
			db.close();
			}
		
		for(HitDatabase db:this.chrom2right.values())
			{
			db.close();
			}
		if(this.rsidSetDb!=null) this.rsidSetDb.close();
		if(this.environment!=null) this.environment.close();
		}
	private static int chrom2int(String s)
		{
		if(s.equals("X")) return 23;
		if(s.equals("Y")) return 24;
		if(s.equals("M")) return 25;
		return Integer.parseInt(s);
		}
	
	private static String int2chrom(int i)
		{
		switch(i)
			{
			case 23: return "X";
			case 24: return "Y";
			case 25: return "M";
			default: return String.valueOf(i);
			}
		}
	
	private void parse(BufferedReader in) throws IOException
		{
		HitBinding binding=new HitBinding();
		int nLine=0;
		byte buffer[]=new byte[HIT_SIZEOF];
		String token[];
		String line=null;
		
		while((line=in.readLine())!=null)
			{
			nLine++;
			if(nLine%100000==0)
				{
				System.err.println("lines:"+nLine+" ..."+line);
				}
			token=TAB.split(line);
			Hit h=new Hit();
			
			h.length= Integer.parseInt(token[4]);
			if(h.length!=expect_flanking_length)
				{
				continue;
				}
			
			char side=token[0].charAt(0);
			assert(side=='L' || side=='R');
			h.rs_id= Integer.parseInt(token[0].substring(1));
			this.rsidSetDb.put(null, h.rs_id, true);
			//this.min_rs_id=Math.min(h.rs_id, this.min_rs_id);
			//this.max_rs_id=Math.max(h.rs_id, this.max_rs_id);
			
			h.strand= token[1].charAt(0);
			assert(h.strand=='+' || h.strand=='-');
			h.chrom_id=chrom2int(token[2]);
			h.position= Integer.parseInt(token[3]);
			
			
			if(token.length>5)
				{
				h.mismatch=token[5];
				}
			
			FileOutputStream db=null;
			if(side=='L')
				{
				HitDatabase hdb=this.chrom2left.get(h.chrom_id);
				if(hdb==null)
					{
					hdb=new HitDatabase(this.base,SUFFIX_L);
					this.chrom2left.put(h.chrom_id, hdb);
					hdb.out=new FileOutputStream(hdb.file);
					}
				db=hdb.out;
				}
			else
				{
				HitDatabase hdb=this.chrom2right.get(h.chrom_id);
				if(hdb==null)
					{
					hdb=new HitDatabase(this.base,SUFFIX_R);
					this.chrom2right.put(h.chrom_id, hdb);
					hdb.out=new FileOutputStream(hdb.file);
					}
				db=hdb.out;
				}
			
			binding.write(h, buffer);
			db.write(buffer);

			}
		for(HitDatabase db: chrom2left.values())
			{
			db.out.flush();
			db.out.close();
			}
		for(HitDatabase db: chrom2right.values())
			{
			db.out.flush();
			db.out.close();
			}
		}
	
	private void sort() throws IOException
		{
		for(HitDatabase db:this.chrom2left.values())
			{
			System.err.println("Sort left "+db.file);
			db.sort();
			}
		System.err.println("Sort right");
		for(HitDatabase db:this.chrom2right.values())
			{
			System.err.println("Sort right "+db.file);
			db.sort();
			}
		}
	
	private void dump() throws IOException
		{
		for(HitDatabase db:this.chrom2left.values())
			{
			db.in=new RandomAccessFile(db.file, "r");
			}
		
		for(HitDatabase db:this.chrom2right.values())
			{
			db.in=new RandomAccessFile(db.file, "r");
			}
		
		Walker<Integer, Boolean> w=this.rsidSetDb.openWalker(null);
		while(w.next())
			{
			int current_rs_id=w.getKey();
			dumpId(current_rs_id);
			}
		
		for(HitDatabase db:this.chrom2left.values())
			{
			db.in.close();
			}
		
		for(HitDatabase db:this.chrom2right.values())
			{
			db.in.close();
			}
		w.close();
		}
	
	private void dumpById(String idFile)
		throws IOException
		{
		for(HitDatabase db:this.chrom2left.values())
			{
			db.in=new RandomAccessFile(db.file, "r");
			}
		
		for(HitDatabase db:this.chrom2right.values())
			{
			db.in=new RandomAccessFile(db.file, "r");
			}
		String line;
		BufferedReader r= IOUtils.openReader(idFile);
		while((line=r.readLine())!=null)
			{
			if(line.trim().isEmpty() || line.startsWith("#")) continue;
			int rs_id=Integer.parseInt(line);
			dumpId(rs_id);
			}
		r.close();
		for(HitDatabase db:this.chrom2left.values())
			{
			db.in.close();
			}
		
		for(HitDatabase db:this.chrom2right.values())
			{
			db.in.close();
			}
		}
	
	private void dumpId(int current_rs_id) throws IOException
		{
		List<Hit> leftHit=new ArrayList<Hit>();
		List<Hit> rightHit=new ArrayList<Hit>();
		
		for(HitDatabase db:this.chrom2left.values())
			{
			leftHit.addAll(equal_range(db.in, current_rs_id));
			}
		for(HitDatabase db:this.chrom2right.values())
			{
			rightHit.addAll(equal_range(db.in, current_rs_id));
			}
		if(leftHit.isEmpty() && rightHit.isEmpty()) return;
		
		
		System.out.println("#4\trs"+current_rs_id+"\t"+leftHit.size()+"\t"+rightHit.size());
		if(leftHit.isEmpty())
			{
			alone(rightHit, 'R', current_rs_id);
			return;
			}
		if(rightHit.isEmpty())
			{
			alone(leftHit, 'L', current_rs_id);
			return;
			}
		challenge(leftHit, rightHit,current_rs_id);
		}
	
	
	private void alone(
			List<Hit> hits,
			char side,
			int rs_id
			)
		{
		for(Hit h:hits)
			{
			System.out.println("#3\trs"+rs_id+"\t"+
				side+"\t"+
				int2chrom(h.chrom_id)+"\t"+
				h.strand+"\t"+
				h.position+"\t"+
				h.length+"\t"+
				h.mismatch
				);
			}
		}
	
	private void push_candidate(CandidatePair candidate,List<CandidatePair> candidates)
		{
		if(candidates.isEmpty() || !strict_use_matches)
			{
			candidates.add(candidate);
			return;
			}
			
		int cmp= candidates.get(0).compareTo( candidate);
		if(cmp > 0)
			{
			candidates.clear();
			candidates.add(candidate);
			return;
			}
		else if(cmp<0)
			{
			return;
			}
		
		candidates.add(candidate);
		Collections.sort(candidates);
		}
	
	private void challenge(
			List<Hit> leftHit,
			List<Hit> rightHit,
			int rs_id
		)
		{
		if(leftHit.isEmpty() && !rightHit.isEmpty())
			{
			alone(rightHit,'R',rs_id);
			return;
			}
		else if(!leftHit.isEmpty() && rightHit.isEmpty())
			{
			alone(leftHit,'L',rs_id);
			return;
			}
			
		List<CandidatePair> candidates= new ArrayList<CandidatePair>();
		for(int i=0;i< leftHit.size();++i)
			{
			Hit L= leftHit.get(i);
			for(int j=0;j< rightHit.size();++j)
				{
				Hit R= rightHit.get(j);
				if(R.strand!=L.strand) continue;
				if(R.chrom_id!=L.chrom_id) continue;
				if(R.strand=='+')
					{
					if(R.position < L.position+L.length) continue;
					if(R.position - (L.position+L.length) > max_len_deletion) continue;
					CandidatePair candidate= new CandidatePair(L,R);
					push_candidate(candidate,candidates);
					}
				else if(R.strand=='-')
					{
					if(L.position < R.position+R.length) continue;
					if(L.position - (R.position+R.length) > max_len_deletion) continue;
					CandidatePair candidate= new CandidatePair(L,R);
					push_candidate(candidate,candidates);
					}
				}
			}
		
		if(candidates.isEmpty())
			{
			System.out.println("#0\trs"+rs_id);
			}
		
		for(CandidatePair candidate:candidates)
			{
			System.out.print(candidates.size()==1?"#1\t":"#2\t");
			candidate.print(System.out);
			System.out.println();
			}
		}
	
	@Override
	protected int processArg(String[] args, int optind)
		{
		if(args[optind].equals("-base"))
			{
			this.base=new File(args[++optind]);
			return optind;
			}
		else if(args[optind].equals("-p"))
			{
			this.program= args[++optind];
			return optind;
			}
		else if(args[optind].equals("-a"))
			{
			this.strict_use_matches=false;
			return optind;
			}
		else if(args[optind].equals("-f"))
			{
			this.ids_file= args[++optind];
			return optind;
			}
		return super.processArg(args, optind);
		}
	
	public static void main(String[] args)
		{
		DigestBowtie2 app=null;
		try
			{
			app=new DigestBowtie2();
			int optind= app.processArgs(args);
			if(app.base==null)
				{
				System.err.println("base dir undefined");
				return;
				}
			
			if(app.program==null)
				{
				System.err.println("undefined program");
				return;
				}
			else if(app.program.equals("sort"))
				{
				if(optind==args.length)
					{
					System.err.println("No arguments");
					return;
					}
				app.open(); 
	            while(optind< args.length)
	                    {
	                  	BufferedReader r=IOUtils.openReader(args[optind++]);
	                    app.parse(r);
	                    r.close();
	                    }
	            app.sort();
				}
			else  if(app.program.equals("dump"))
				{
				app.open();
				if(app.ids_file==null)
					{
					app.dump();
					}
				else
					{
					app.dumpById(app.ids_file);
					}
				}
			else
				{
				System.err.println("unknown program");
				}
			}
		catch (Exception e)
			{
			e.printStackTrace();
			}
		finally
			{
			if(app!=null) app.close();
			}
		}
	}
